package main

import (
	"archive/zip"
	"bytes"
	"crypto/rand"
	"embed"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/sync/errgroup"
)

//go:embed website
var website embed.FS

const PlatformKEY = "NSAWFJHUROFWKFJF"

func main() {
	f, err := fs.Sub(website, "website")
	if err != nil {
		log.Fatalf("fs.Sub error: %v", err)
	}

	mux := http.NewServeMux()
	mux.Handle("/", http.FileServer(http.FS(f)))
	mux.HandleFunc("/encrypt", encrypt)
	mux.HandleFunc("/decrypt", decrypt)

	server := http.Server{
		Addr:         "0.0.0.0:80",
		Handler:      mux,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 10 * time.Second,
	}

	fmt.Printf("Starting server at %s...\n", server.Addr)
	server.ListenAndServe()
}

func encrypt(w http.ResponseWriter, r *http.Request) {
	// 读取表单
	public := r.FormValue("public")
	publicPath, err := WriteTempFile([]byte(public), "public-*.pem")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	file, hdr, err := r.FormFile("key")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer file.Close()

	// 直接提供下载文件给用户
	name := strings.Replace(hdr.Filename, filepath.Ext(hdr.Filename), ".zip", -1)
	w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
	w.Header().Set("Content-Disposition", "attachment; filename="+name)

	zw := zip.NewWriter(w)
	if err := zw.SetComment(hdr.Filename); err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer zw.Close()

	// 返回给用户的压缩包
	dek, err := zw.Create("dek.cipher")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	data, err := zw.Create("data.cipher")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	mac, err := zw.Create("hmac.bin")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	// 使用公钥加密随机生成的 DEK 保存到压缩包的 dek.cipher 文件中
	key := RandomHexKey(16)
	stderr := bytes.Buffer{}
	cmd := exec.CommandContext(r.Context(), "tongsuo", "pkeyutl", "-inkey", publicPath, "-pubin")
	cmd.Stdin = strings.NewReader(key)
	cmd.Stdout = dek
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		http.Error(w, fmt.Sprintf("stderr: %v error: %v", stderr.String(), err), 500)
		return
	}

	// 读取用户上传的文件，加密同时对原文使用平台密钥进行 HMAC，在校验 HMAC 时候证明是文件是本平台加密的
	wg := errgroup.Group{}

	pr, pw := io.Pipe()
	source := io.TeeReader(file, pw)

	// 使用 SM3 对数据进行 HMAC
	wg.Go(func() error {
		stderr := bytes.Buffer{}
		cmd := exec.CommandContext(r.Context(), "tongsuo", "dgst", "-hmac", PlatformKEY, "-sm3", "-binary")
		cmd.Stdout = mac
		cmd.Stdin = pr
		cmd.Stderr = &stderr

		if err := cmd.Run(); err != nil {
			return fmt.Errorf("cmd.Run stderr: %s error: %v", stderr.String(), err)
		}

		return nil
	})

	// 使用 SM4 和随机生成的 DEK 对数据进行加密
	wg.Go(func() error {
		stderr := bytes.Buffer{}
		cmd := exec.CommandContext(r.Context(), "tongsuo", "enc", "-K", key, "-e", "-sm4-ofb", "-iv", key, "-nosalt")
		cmd.Stdin = source
		cmd.Stdout = data
		cmd.Stderr = &stderr

		if err := cmd.Run(); err != nil {
			return fmt.Errorf("cmd.Run stderr: %s error: %v", stderr.String(), err)
		}

		return nil
	})

	if err := wg.Wait(); err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
}

func decrypt(w http.ResponseWriter, r *http.Request) {
	// 读取表单
	private := r.FormValue("private")
	privatePath, err := WriteTempFile([]byte(private), "private-*.pem")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	file, _, err := r.FormFile("key")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer file.Close()

	fileSize, err := file.Seek(0, 2)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	_, err = file.Seek(0, 0)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	// 直接提供下载文件给用户
	zr, err := zip.NewReader(file, fileSize)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
	w.Header().Set("Content-Disposition", "attachment; filename="+zr.Comment)

	dek, err := zr.Open("dek.cipher")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer dek.Close()

	mac, err := zr.Open("hmac.bin")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer mac.Close()

	data, err := zr.Open("data.cipher")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer data.Close()

	// 首先读取 dek 用私钥解密这个 DEK
	dekPlain := bytes.Buffer{}
	stderr := bytes.Buffer{}
	cmd := exec.CommandContext(r.Context(), "tongsuo", "pkeyutl", "-decrypt", "-inkey", privatePath)
	cmd.Stdout = &dekPlain
	cmd.Stdin = dek
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		http.Error(w, fmt.Sprintf("stderr: %v error: %v", stderr.String(), err), 500)
		return
	}

	wg := errgroup.Group{}

	pr, pw := io.Pipe()

	// 使用 SM3 对数据进行 HMAC
	wg.Go(func() error {
		macBin := bytes.Buffer{}
		stderr := bytes.Buffer{}
		cmd := exec.CommandContext(r.Context(), "tongsuo", "dgst", "-hmac", PlatformKEY, "-sm3", "-binary")
		cmd.Stdout = &macBin
		cmd.Stdin = pr
		cmd.Stderr = &stderr

		if err := cmd.Run(); err != nil {
			return fmt.Errorf("cmd.Run stderr: %s error: %v", stderr.String(), err)
		}

		macOrigin, err := io.ReadAll(mac)
		if err != nil {
			return err
		}

		ok := bytes.Equal(macBin.Bytes(), macOrigin)
		if !ok {
			return errors.New("hmac verification failed")
		}

		return nil
	})

	// 用 DEK 解密文件
	wg.Go(func() error {
		stderr := bytes.Buffer{}
		cmd = exec.CommandContext(r.Context(), "tongsuo", "enc", "-K", dekPlain.String(), "-d", "-sm4-ofb", "-iv", dekPlain.String(), "-nosalt")
		cmd.Stdout = pw
		cmd.Stdin = data
		cmd.Stderr = &stderr

		if err := cmd.Run(); err != nil {
			return fmt.Errorf("cmd.Run stderr: %s error: %v", stderr.String(), err)
		}

		return nil
	})

	if err := wg.Wait(); err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	// 校验通过之后，重新打开文件，解密并传递给用户
	if err := data.Close(); err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	stderr.Reset()
	cmd = exec.CommandContext(r.Context(), "tongsuo", "enc", "-K", dekPlain.String(), "-d", "-sm4-ofb", "-iv", dekPlain.String(), "-nosalt")
	cmd.Stdout = w
	cmd.Stdin = data
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		http.Error(w, fmt.Sprintf("stderr: %v error: %v", stderr.String(), err), 500)
		return
	}
}

func RandomHexKey(length int) string {
	k := make([]byte, length)
	io.ReadFull(rand.Reader, k)
	return hex.EncodeToString(k)
}

func WriteTempFile(data []byte, name string) (string, error) {
	if len(data) < 1 {
		return "", io.ErrShortWrite
	}

	f, err := os.CreateTemp(os.TempDir(), name)
	if err != nil {
		return "", err
	}

	n, err := f.Write(data)
	if err != nil {
		return "", err
	}

	if n != len(data) {
		return "", io.ErrUnexpectedEOF
	}

	return f.Name(), nil
}
